{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "810ce279",
   "metadata": {},
   "source": [
    "# Document-level access example using the push document APIs\n",
    "\n",
    "In Azure AI Search, you can upload any JSON document payload to a search index for indexing. This notebook shows you how index documents that contain [user access permissions at the document level](azure/search/search-document-level-access-overview), and then query the index to return only those results that the user is authorized to view.\n",
    "\n",
    "The security principal behind the query access token determines the \"user\". The permission metadata in the document determines whether the user has authorization to the content. Internally, the search engine filters out any documents that aren't associated with the security principal.\n",
    "\n",
    "This feature is currently in preview.\n",
    "\n",
    "For an alternative approaching using indexers and pull API, see [Quickstart-Document-Permissions-Pull-API](../Quickstart-Document-Permissions-Pull-API/document-permissions-pull-api.ipynb).\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6585426",
   "metadata": {},
   "source": [
    "## Prerequisites\n",
    "\n",
    "+ Azure AI Search, with [role-based access control](https://learn.microsoft.com/azure/search/search-security-enable-roles).\n",
    "\n",
    "## Permissions\n",
    "\n",
    "This walkthrough uses Microsoft Entra ID authentication and authorization.\n",
    "\n",
    "On Azure AI Search, you must have role assignments to create objects and run queries:\n",
    "\n",
    "+ **Search Service Contributor**\n",
    "+ **Search Index Data Contributor**\n",
    "+ **Search Index Data Reader**\n",
    "\n",
    "For more information, see [Connect to Azure AI Search using roles](https://learn.microsoft.com/azure/search/search-security-rbac) and [Quickstart: Connect without keys for local testing](https://learn.microsoft.com/azure/search/search-get-started-rbac).\n",
    "\n",
    "## Set the environment variables\n",
    "\n",
    "1. Rename `sample.env` to `.env`.\n",
    "1. In the `.env` file, provide a full endpoint to your search service (https://your-search-service.search.windows.net).\n",
    "1. Replace the default index name if you want a different name.\n",
    "\n",
    "## Load Connections\n",
    "\n",
    "We recommend creating a virtual environment to run this sample code. In Visual Studio Code, open the control palette (ctrl-shift-p) to create an environment. This notebook was tested on Python 3.10.\n",
    "\n",
    "Once your environment is created, load the environment variables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2975a7f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dotenv import load_dotenv\n",
    "from azure.identity import DefaultAzureCredential, get_bearer_token_provider\n",
    "import os\n",
    "\n",
    "load_dotenv(override=True) # take environment variables from .env.\n",
    "\n",
    "# The following variables from your .env file are used in this notebook\n",
    "endpoint = os.environ[\"AZURE_SEARCH_ENDPOINT\"]\n",
    "credential = DefaultAzureCredential()\n",
    "index_name = os.getenv(\"AZURE_SEARCH_INDEX\")\n",
    "token_provider = get_bearer_token_provider(credential, \"https://search.azure.com/.default\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9327cf01",
   "metadata": {},
   "source": [
    "## Create Sample Index\n",
    "\n",
    "The search index must includes fields for your content and for permission metadata. Assign the new permission filter option to a string field and make sure the field is filterable. The search engine builds the filter internally at query time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9863061f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.search.documents.indexes.models import SearchField, SearchIndex, PermissionFilter, SearchIndexPermissionFilterOption\n",
    "from azure.search.documents.indexes import SearchIndexClient\n",
    "\n",
    "index_client = SearchIndexClient(endpoint=endpoint, credential=credential)\n",
    "index = SearchIndex(\n",
    "    name=index_name,\n",
    "    fields=[\n",
    "        SearchField(name=\"id\", type=\"Edm.String\", key=True, filterable=True, sortable=True),\n",
    "        SearchField(name=\"oid\", type=\"Collection(Edm.String)\", retrievable=True, filterable=True, permission_filter=PermissionFilter.USER_IDS),\n",
    "        SearchField(name=\"group\", type=\"Collection(Edm.String)\", retrievable=True, filterable=True, permission_filter=PermissionFilter.GROUP_IDS),\n",
    "        SearchField(name=\"name\", type=\"Edm.String\", searchable=True)\n",
    "    ],\n",
    "    permission_filter_option=SearchIndexPermissionFilterOption.ENABLED\n",
    ")\n",
    "\n",
    "index_client.create_index(index=index)\n",
    "print(f\"Index '{index_name}' created with permission filter option enabled.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5cf4169",
   "metadata": {},
   "source": [
    "## Connect to Graph to find your object ID (OID) and groups\n",
    "\n",
    "This step calls the Graph APIs to get a few group IDs for your Microsoft Entra identity. Your group IDs will be added to the access control list of the objects created in the next step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63904f09",
   "metadata": {},
   "outputs": [],
   "source": [
    "from msgraph import GraphServiceClient\n",
    "client = GraphServiceClient(credentials=credential, scopes=[\"https://graph.microsoft.com/.default\"])\n",
    "\n",
    "groups = await client.me.member_of.get()\n",
    "me = await client.me.get()\n",
    "oid = me.id"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9ce6d0f",
   "metadata": {},
   "source": [
    "## Upload Sample Data\n",
    "\n",
    "This step creates the container, folders, and uploads documents into Azure Storage. It assigns your group IDs to to the access control list for each file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8fb830a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.search.documents import SearchClient\n",
    "search_client = SearchClient(endpoint=endpoint, index_name=index_name, credential=credential)\n",
    "\n",
    "documents = [\n",
    "    { \"id\": \"1\", \"oid\": [oid], \"group\": [groups.value[0].id], \"name\": \"Document 1\" },\n",
    "    { \"id\": \"2\", \"oid\": [\"all\"], \"group\": [groups.value[0].id], \"name\": \"Document 2\" },\n",
    "    { \"id\": \"3\", \"oid\": [oid], \"group\": [\"all\"], \"name\": \"Document 3\" },\n",
    "    { \"id\": \"4\", \"oid\": [\"none\"], \"group\": [\"none\"], \"name\": \"Document 4\" },\n",
    "    { \"id\": \"5\", \"oid\": [\"none\"], \"group\": [groups.value[0].id], \"name\": \"Document 5\" },\n",
    "]\n",
    "search_client.upload_documents(documents=documents)\n",
    "print(\"Documents uploaded to the index.\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5c93f76",
   "metadata": {},
   "source": [
    "## Search sample data with x-ms-query-source-authorization\n",
    "\n",
    "This query uses an empty search string (`*`) to provide an unqualified search. It returns the file name and permission metadata associated with each file. Notice that each file is associated with a different group ID."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd872e8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = search_client.search(search_text=\"*\", x_ms_query_source_authorization=token_provider(), select=\"name,oid,group\", order_by=\"id asc\")\n",
    "\n",
    "for result in results:\n",
    "    print(f\"Name: {result['name']}, OID: {result['oid']}, Group: {result['group']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d31b67d8",
   "metadata": {},
   "source": [
    "## Search sample data without x-ms-query-source-authorization \n",
    "\n",
    "This step demonstrates the user experience when authorization fails. No results are returned in the response."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1f2f2a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = search_client.search(search_text=\"*\", x_ms_query_source_authorization=None, select=\"name,oid,group\", order_by=\"id asc\")\n",
    "\n",
    "for result in results:\n",
    "    print(f\"Name: {result['name']}, OID: {result['oid']}, Group: {result['group']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ad253ec",
   "metadata": {},
   "source": [
    "## Next steps\n",
    "\n",
    "To learn more, see [Document-level access control in Azure AI Search](https://learn.microsoft.com/azure/search/search-document-level-access-overview)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
